# 创建数据库引擎
from contextlib import contextmanager, asynccontextmanager
from typing import Annotated
from fastapi import Depends
from sqlalchemy import event, text, column
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from sqlalchemy.orm import ORMExecuteState, sessionmaker
from sqlmodel import SQLModel, create_engine, Session
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

from src.base.context import tenant_var, data_scope_var
# asyncmy 库对BIGINT数据类型支持不完善，
# 由于使用 Cython 加速，性能很强，但是需要C扩展,调试困难，近两年更新少
# MYSQL_URL = "mysql+asyncmy://root:123456@127.0.0.1/fin_tax_analysis?charset=utf8mb4"
MYSQL_URL = "mysql+aiomysql://root:123456@127.0.0.1/fin_tax_analysis?charset=utf8mb4"
# # 1. 创建数据库引擎,create_engine 同步引擎
# engine = create_engine(
#     MYSQL_URL,
#     pool_size=20,  # 连接池大小
#     max_overflow=10,  # 最大溢出连接数
#     pool_recycle=3600,  # 连接回收时间(秒)
#     pool_pre_ping=True,  # 执行前健康检查
#     echo=True,  # SQL日志
# )
# # 2. 创建会话工厂，并把当前引擎绑定给该对话
# SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False, expire_on_commit=True)
#
# # 3. 依赖生成器（FastAPI推荐方式）
# @contextmanager
# def db_session() -> Session:
#     """创建数据库会话上下文管理器"""
#     with SessionLocal() as session:
#         try:
#             yield session
#             session.commit()  # 自动提交成功操作
#         except Exception:
#             session.rollback()  # 异常回滚
#             raise
#
# dbClient = Session(engine)
# def create_db_and_tables():
#     SQLModel.metadata.create_all(engine)

async_engine = create_async_engine(
    MYSQL_URL,
    pool_size=20,  # 连接池大小
    max_overflow=10,  # 最大溢出连接数
    pool_recycle=3600,  # 连接回收时间(秒)
    pool_pre_ping=True,  # 执行前健康检查
    echo=True,  # SQL日志
)

AsyncSessionLocal = async_sessionmaker(bind=async_engine, autocommit=False, autoflush=False, expire_on_commit=True,
                                       class_=AsyncSession)

#  @asynccontextmanager 让你用一个简单的异步函数，就能写出安全、可重用的 async with 资源管理代码，是异步编程中的“语法糖”利器。
@asynccontextmanager
async def async_db_session() -> AsyncSession:
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()  # 自动提交成功操作
        except Exception:
            await session.rollback()  # 异常回滚
            raise


asyncDbClient = AsyncSession(async_engine)
async def get_db() -> AsyncSession:
    async with async_db_session() as session:
        yield session


DbSession = Annotated[AsyncSession, Depends(get_db)]


@event.listens_for(Session, "do_orm_execute", retval=True)
def apply_dept_filter(execute_state: ORMExecuteState):
    if not execute_state.is_select:
        return
    data_scope_sql_handle(execute_state)
    tenant_sql_handle(execute_state)


def tenant_sql_handle(execute_state):
    is_tenant = hasattr(execute_state.bind_mapper.class_, "tenant_id")
    tenant_dict = tenant_var.get()
    if tenant_dict and is_tenant:
        tenant_table_alias = tenant_dict['tenant_table_alias']
        tenant_id_alias = tenant_dict['tenant_id_alias']
        tenant_id_value = tenant_dict['tenant_id_value']
        execute_state.statement = execute_state.statement.where(
            text(f"{tenant_table_alias}.{tenant_id_alias} = {tenant_id_value}"))


def data_scope_sql_handle(execute_state):
    data_scope_dict = data_scope_var.get()
    if data_scope_dict:
        dept_table_alias = data_scope_dict['dept_table_alias']
        dept_id_alias = data_scope_dict['dept_id_alias']
        user_table_alias = data_scope_dict['user_table_alias']
        user_id_alias = data_scope_dict['user_id_alias']
        user_id_values = data_scope_dict["user_id_values"]
        dept_id_values = data_scope_dict["dept_id_values"]

        dept_where = None
        user_where = None
        if dept_id_values and len(dept_id_values) > 0:
            dept_where = f"{dept_table_alias}.{dept_id_alias} in ({', '.join(map(str, dept_id_values))})"
        if user_id_values and len(user_id_values) > 0:
            user_where = f"{user_table_alias}.{user_id_alias} in ({', '.join(map(str, user_id_values))})"
        data_scope_sql = None
        if dept_where and user_where:
            data_scope_sql = f" ( {dept_where}  OR {user_where} )"
        elif dept_where:
            data_scope_sql = f" {dept_where}"
        elif user_where:
            data_scope_sql = f" {user_where}"
        if data_scope_sql:
            execute_state.statement = execute_state.statement.where(text(data_scope_sql))


class DbMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        with async_db_session() as dbSession:
            try:
                request.state.dbSession = dbSession  # 挂接到请求对象
                response = await call_next(request)
            except Exception:
                dbSession.rollback()
                raise
        return response


# ================== redis ===============================

import redis

REDIS_HOST = '127.0.0.1'
REDIS_PORT = 6379

# 创建一个全局的Redis连接池
redis_pool = redis.ConnectionPool(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
redisClient = redis.Redis(connection_pool=redis_pool)
